;;===============================================================
;; ahocorasick
;; Implements the Aho-Corasick automaton itself.
;; Uses Trie (in trie.clj) and the ACPointer structure
;; Mainly exports the ac function which receives a list of words
;; and returns a matcher function which receives a text and returns the matches.
;;===============================================================
(ns strclj.ahocorasick
  (:use [strclj.trie]
        [objects.node])
  (:refer strclj.trie)
  (:import [objects Node])
)

;;===============================================================
;; ACPointer - wraps a Node object and a long int specifying the prefix length.
;; Useful for building automaton links.
;;===============================================================
;  Pointer protocol: ACPointer will implement it.
(defprotocol Pointer
  (get-node [this])
  (prefix-length [this]))

; Implements the Pointer functions for nil.
(extend-type nil
  Pointer
  (get-node [_] nil)
  (prefix-length [_] 0)
)

; ACPointer implementation. 
; basically a wrapper which contains a Trie's node and its prefix length.
(defrecord ACPointer
  [^Node node ^long prefix-length]
  Pointer
  (get-node [this] node)
  (prefix-length [_] prefix-length)
)

(defn ^ACPointer pointer-make
  "Generates ACPointer for a given node & prefix length."
  [^Node node ^long prefix-length]
  (->ACPointer node prefix-length)
)

(defn- set-node-pointer!
  "Adds a pointer to a given node into a node's dict, for a given key."
  [^Node node ^Node pointed ^long length key]
  (node-set-dict! node 
    (assoc (node-get-dict node) key (pointer-make pointed length))
  )
)

; Setters & getters for the fail & output pointers.
(defn- set-fail-pointer! 
  [^Node node ^Node fail-node ^long length]
  (set-node-pointer! node fail-node length :fail)
)
  
(defn- ^ACPointer get-fail-pointer
  [^Node node]
  (get (node-get-dict node) :fail)
)

(defn- set-output-pointer! 
  [^Node node ^Node output-node ^long length]
  (set-node-pointer! node output-node length :output)
)

(defn- ^ACPointer get-output-pointer
  [^Node node]
  (get (node-get-dict node) :output)
)


(defn- generate-sequence
  "Receives a function and a node and generates a sequence of ACPointers 
  by applying the function recursively, consing each node in the process.
  In other words, it produces the appropriate path from the root up to the node,
  depending on the given function.
  e.g., if the fn follows the fail pointer - we have the full path of the fail links (longest suffixes)
  from the node back to the root."
  [get-function ^Node node]
  (letfn [(generate [^Node current-node]
                    (when-let [chosen (get-function current-node)]
                      (cons chosen (lazy-seq (generate (get-node chosen))))))]
    (generate node))  
)

(defn- generate-fail-seq
  "Generate a sequence of the fail nodes."
  [^Node node]
  (generate-sequence get-fail-pointer node)
)

(defn- generate-output-seq
  "Generate a sequence of the output nodes."
  [^Node node]
  (generate-sequence get-output-pointer node)
)


(defn- ^ACPointer get-longest-suffix
  "Generates the failure node pointer. per AC algorithm, it is the longest proper suffix of the string."
  [^Node root ^Node father key]
  (let [father-fail-seq (generate-fail-seq father)
        longest-suffix-node (first (filter (fn [pointer] (get-child (get-node pointer) key)) father-fail-seq))]
    (if longest-suffix-node 
      longest-suffix-node 
      (when (get-child root key) (pointer-make root 0)))
  )
)


(defn upgrade-trie!
  "Receives a simple Trie and 'upgrades' it to an Aho-Corasick automaton
  by adding failure and output links."
  [^Node root]
  (let [q (java.util.LinkedList.)]
    ; Add all the children of the root's kids to the queue (the fail links of the root's kids are the root itself).
    (doseq [[_ child] (get-children root)]
      (doseq [subchild (get-children child)]
        (. q add [child subchild])))
    ; (while !isEmpty) :)
    (when-not (. q isEmpty)
      (loop [[father [key child]] (. q removeFirst)]
        (when-let [fail-node (get-longest-suffix root father key)]
          (let [fail-node-child (get-child (get-node fail-node) key)
                length (inc (prefix-length fail-node))]
            ; Failure node
            (set-fail-pointer! child fail-node-child length)
            ; Output node
            (cond
              ; Finished word
              (get-finish-flag fail-node-child) 
                (set-output-pointer! child fail-node-child length)
              ; Not finished word
              (get-output-pointer fail-node-child)
	              (let [output-pointer (get-output-pointer fail-node-child)]
	                (set-output-pointer! child output-pointer (prefix-length (get-output-pointer fail-node-child)))
	              )
            )
          )
        )
        (doseq [subchild (get-children child)]
          (. q add [child subchild])
        )
        (when-not (. q isEmpty)
          (recur (. q removeFirst))
        )
      )
    )   
  )
)


(defn- match-node
  "Get list of match results for a given node. 
  Just need to follow the output pointers and map them into indexes (relative to the given index)."
  [^Node node ^long idx]
  ; Define a mapping between pointers and indexes
  (letfn [(pointers-to-indexes 
            [^Node node ^long idx]
            (map (fn [^ACPointer ptr] [(- idx (prefix-length ptr) -1) idx])
                 (generate-output-seq node)
            )
         )]
    (if (get-finish-flag node)
	    (cons [0 idx] (pointers-to-indexes node idx))
	    (pointers-to-indexes node idx)
    )
  )
)


(defn- ac-match-step
  "
  Performs the matching process recursively - each time going from text[i] to text[i+1].
  text: the remaining text
  text-idx: index of the current character
  start-text: the chunk of text in which we start the search, determined by the fail pointers 
  (this is the very advantage of the AC alg. - we don't necessarily need to start searching from the start) 
  start-text-idx: the index of the current char from which we start searching.
  root: the root of the automaton.
  node: the current node.
  This function calculates the next parameters each time, and recursively creates a lazy sequence
  which stores the matches.
  "
  [text text-idx start-text start-text-idx ^Node root ^Node node]
  (if (first text) 
    (let [current-char (first text)
          child (get-child node current-char)
          fail-child (when-not child
                       (let [fail-ptr (get-longest-suffix root node current-char)]
	                       (if fail-ptr 
	                         (pointer-make (get-child (get-node fail-ptr) current-char) (inc (prefix-length fail-ptr)))
	                         (pointer-make root 0)
	                       )
                       )
                     )
          next-start-text-idx (if child start-text-idx (- text-idx (prefix-length fail-child) -1))
          next-start-text (if child start-text (drop (- next-start-text-idx start-text-idx) start-text))
          next-node (if child child (get-node fail-child))
          matches (match-node next-node (- text-idx next-start-text-idx))
          ]
      (lazy-seq
        (if (seq matches)
          (concat (map(fn [item] (map (partial + next-start-text-idx) item)) matches)
                  (ac-match-step (rest text) (inc text-idx) next-start-text next-start-text-idx root next-node)
          )
          (ac-match-step (rest text) (inc text-idx) next-start-text next-start-text-idx root next-node)
        )
      )
    )
  )
)

(defn ac-make
  "Receives a list of words, creates an Aho-Corasick automaton and returns it."
  [words]
  (let [trie (trie-make words)]
    (upgrade-trie! trie)
    trie
  )
)

(defn ac-match
  "Finds all matches for the automaton in the text.
  By calling ac-match-step with the appropriate starting parameters."
  [^Node automaton text]
  (ac-match-step text 0 text 0 automaton automaton)
)

(defn ac-matcher
  "Generates a matcher for a given automaton,
  by creating a partial function using ac-match."
  [^Node automaton] 
  (partial ac-match automaton)  
)

(defn ac 
  [words]
  "This is the headline function: receives a list of words and returns a matcher function."
  (let [automaton (ac-make words)]
    (ac-matcher automaton)
  )
)



